import os
import json
import time
from app.client.BotClient import (
    NewKnowledgeBaseClient,
    KnowledgeBaseClient,
    KnowledgeBaseUploadClient,
    KnowledgeBaseListFiles,
    BotInitClient,
    BotInfoClient,
)

from utils.tools import logger, transform_kb_data_with_pandas, root_dir


class BotManager:
    def __init__(self, knowledge_name, knowledge_base_folder="data"):
        self.knowledge_name = knowledge_name
        self.folder_path = os.path.join(
            os.getcwd(), knowledge_base_folder, self.knowledge_name
        )
        self.knowledge_base_db = self.get_knowledge_base()
        self.bot_name = f"{knowledge_name}_BOT"
        self.kb_ids = []

    def get_knowledge_base(self):
        try:
            kbc = KnowledgeBaseClient()
            return transform_kb_data_with_pandas(kbc.get_knowledge_base()["data"])
        except Exception as e:
            logger.error(f"获取知识库数据失败: {e}")
            return {}

    def check_knowledge_base(self):
        if self.knowledge_name not in self.knowledge_base_db:
            self.create_and_upload_knowledge_base()
        else:
            self.kb_ids = self.knowledge_base_db[self.knowledge_name]

    def create_and_upload_knowledge_base(self):
        try:
            # file_name = f"{self.knowledge_name}.txt"
            # file_path = os.path.join(self.folder_path, self.knowledge_name, file_name)
            # os.makedirs(os.path.dirname(file_path), exist_ok=True)

            nkbc = NewKnowledgeBaseClient(kb_name=self.knowledge_name)
            result = nkbc.new_knowledge_base()
            logger.info(f"【{self.knowledge_name}】 create knowledge_base susccess")
            logger.info(result.get("msg"))
            self.kb_id = result.get("data", {}).get("kb_id")
            self.kb_name = result.get("data", {}).get("kb_name")

            kbuc = KnowledgeBaseUploadClient(self.folder_path, self.kb_id)
            kbuc.upload_files()
            logger.info(f"【{self.knowledge_name}】 upload files susccess")

            self.kb_ids = [self.kb_id]
        except Exception as e:
            logger.error(f"创建和上传知识库失败: {e}")

    def check_files_status(self):
        for kb_id in self.kb_ids:
            try:
                kblf = KnowledgeBaseListFiles(kb_id=kb_id)

                start_time = time.time()
                while True:
                    files_detail = (
                        kblf.get_list_files().get("data", {}).get("details", [])
                    )
                    logger.info(
                        f"【{self.knowledge_name}】 下有{len(files_detail)}个文件"
                    )
                    all_green = all(
                        file_detail.get("status") == "green"
                        for file_detail in files_detail
                    )
                    if all_green or (time.time() - start_time) > 300:  # 超时设置为5分钟
                        break
                    time.sleep(1)
                    logger.warning(
                        f"【{self.knowledge_name}】文件解析未完成，请稍等..."
                    )

                if all_green:
                    logger.info(f"【{self.knowledge_name}】文件解析完成")
                else:
                    logger.error(f"文件解析超时失败")
            except Exception as e:
                logger.error(f"检查文件状态失败: {e}")

    def init_bot_if_not_exists(self):
        try:
            bot_info_list = BotInfoClient().get_bot_info().get("data", [])
            for bot_info in bot_info_list:
                if bot_info.get("bot_name") == self.bot_name:
                    logger.warning(f"机器人【{self.bot_name}】已存在，将不创建")
                    return bot_info

            bic = BotInitClient(self.kb_ids, self.bot_name)
            bic.init_bot()
            logger.info(f"机器人【{self.bot_name}】创建成功")

            return None
        except Exception as e:
            logger.error(f"初始化机器人失败: {e}")

    def run(self):
        try:
            self.check_knowledge_base()
            self.check_files_status()
            bot_info = self.init_bot_if_not_exists()

            if not bot_info:
                bot_info = self.init_bot_if_not_exists()
                logger.info("获取机器人信息")
                logger.info(json.dumps(bot_info, ensure_ascii=False, indent=4))

            self.save_bot_info(bot_info)
            logger.info(
                f"{self.bot_name}机器人配置信息下载成功，请读取配置文件进行后续操作"
            )
            return True
        except Exception as e:
            logger.error(f"运行过程中出现错误: {e}")
            raise Exception(f"运行过程中出现错误: {e}")

    def save_bot_info(self, bot_info):
        try:
            with open(
                os.path.join(root_dir, "bot_conf", f"{self.bot_name}.json"),
                "w",
                encoding="utf-8",
            ) as f:
                f.write(json.dumps(bot_info, ensure_ascii=False, indent=4))
        except Exception as e:
            logger.error(f"保存机器人信息失败: {e}")


if __name__ == "__main__":
    knowledge_name = "心理学经典案例揭秘：都市变态心理解析"
    manager = BotManager(knowledge_name)
    manager.run()
